#include "stdafx.h"
#include "net/OAuth10Credentials.hpp"
#include "net/HTTPAuthenticationParams.hpp"
#include "stdlib/exception.hpp"
#include "stdlib/timestamp.hpp"
using JHCPP::stdlib::CTimeStamp;
using JHCPP::stdlib::jhException::NotAuthenticatedException;
#include "utils/format.hpp"
#include "utils/numberformatter.hpp"
using JHCPP::utils::format;
using JHCPP::utils::CNumberFormatter;
#include "stream/base64encode.hpp"
using JHCPP::stream::CBase64Encode;
#include "math/randomstream.hpp"
using JHCPP::math::CRandomInputStream;
#include "code/hmacengine.hpp"
#include "code/sha1engine.hpp"
#include "code/digestengine.hpp"
using namespace JHCPP::code;
#include <map>
#include <sstream>


NET_NAMESPACE_BEGIN

const std::string COAuth10Credentials::SCHEME = "OAuth";

COAuth10Credentials::COAuth10Credentials()
{
}

COAuth10Credentials::COAuth10Credentials(const std::string& consumerKey, const std::string& consumerSecret)
	: m_consumerKey(consumerKey),
	m_consumerSecret(consumerSecret)
{
}

COAuth10Credentials::COAuth10Credentials(const std::string& consumerKey, const std::string& consumerSecret, const std::string& token, const std::string& tokenSecret)
	: m_consumerKey(consumerKey),
	m_consumerSecret(consumerSecret),
	m_token(token),
	m_tokenSecret(tokenSecret)
{
}

COAuth10Credentials::COAuth10Credentials(const CHTTPRequest& request)
{
	if (request.hasCredentials())
	{
		std::string authScheme;
		std::string authParams;
		request.getCredentials(authScheme, authParams);
		if (icompare(authScheme, SCHEME) == 0)
		{
			CHTTPAuthenticationParams params(authParams);
			std::string consumerKey = params.get("oauth_consumer_key", "");
			CURI::decode(consumerKey, m_consumerKey);
			std::string token = params.get("oauth_token", "");
			CURI::decode(token, m_token);
			std::string callback = params.get("oauth_callback", "");
			CURI::decode(callback, m_callback);
		}
		else throw NotAuthenticatedException("No OAuth credentials in Authorization header", authScheme);
	}
	else throw NotAuthenticatedException("No Authorization header found");
}

COAuth10Credentials::~COAuth10Credentials()
{
}

void COAuth10Credentials::setConsumerKey(const std::string& consumerKey)
{
	m_consumerKey = consumerKey;
}

const std::string& COAuth10Credentials::getConsumerKey() const
{
	return m_consumerKey;
}

void COAuth10Credentials::setConsumerSecret(const std::string& consumerSecret)
{
	m_consumerSecret = consumerSecret;
}

const std::string& COAuth10Credentials::getConsumerSecret() const
{
	return m_consumerSecret;
}

void COAuth10Credentials::setToken(const std::string& token)
{
	m_token = token;
}

const std::string& COAuth10Credentials::getToken() const
{
	return m_token;
}

void COAuth10Credentials::setTokenSecret(const std::string& tokenSecret)
{
	m_tokenSecret = tokenSecret;
}

const std::string& COAuth10Credentials::getTokenSecret() const
{
	return m_tokenSecret;
}

void COAuth10Credentials::setRealm(const std::string& realm)
{
	m_realm = realm;
}

const std::string& COAuth10Credentials::getRealm() const
{
	return m_realm;
}

void COAuth10Credentials::setCallback(const std::string& callback)
{
	m_callback = callback;
}

const std::string& COAuth10Credentials::getCallback() const
{
	return m_callback;
}

void COAuth10Credentials::authenticate(CHTTPRequest& request, const CURI& uri, SignatureMethod method)
{
	CHTMLForm emptyParams;
	authenticate(request, uri, emptyParams, method);
}

void COAuth10Credentials::authenticate(CHTTPRequest& request, const CURI& uri, const CHTMLForm& params, SignatureMethod method)
{
	if (method == SIGN_PLAINTEXT)
	{
		signPlaintext(request);
	}
	else
	{
		CURI uriWithoutQuery(uri);
		uriWithoutQuery.setQuery("");
		uriWithoutQuery.setFragment("");
		signHMACSHA1(request, uriWithoutQuery.toString(), params);
	}
}

bool COAuth10Credentials::verify(const CHTTPRequest& request, const CURI& uri)
{
	CHTMLForm params;
	return verify(request, uri, params);
}

bool COAuth10Credentials::verify(const CHTTPRequest& request, const CURI& uri, const CHTMLForm& params)
{
	if (request.hasCredentials())
	{
		std::string authScheme;
		std::string authParams;
		request.getCredentials(authScheme, authParams);
		if (icompare(authScheme, SCHEME) == 0)
		{
			CHTTPAuthenticationParams oauthParams(authParams);

			std::string version = oauthParams.get("oauth_version", "1.0");
			if (version != "1.0") throw NotAuthenticatedException("Unsupported OAuth version", version);

			m_consumerKey.clear();
			std::string consumerKey = oauthParams.get("oauth_consumer_key", "");
			CURI::decode(consumerKey, m_consumerKey);

			m_token.clear();
			std::string token = oauthParams.get("oauth_token", "");
			CURI::decode(token, m_token);

			m_callback.clear();
			std::string callback = oauthParams.get("oauth_callback", "");
			CURI::decode(callback, m_callback);

			std::string nonceEnc = oauthParams.get("oauth_nonce", "");
			std::string nonce;
			CURI::decode(nonceEnc, nonce);

			std::string timestamp = oauthParams.get("oauth_timestamp", "");

			std::string method = oauthParams.get("oauth_signature_method", "");

			std::string signatureEnc = oauthParams.get("oauth_signature", "");
			std::string signature;
			CURI::decode(signatureEnc, signature);

			std::string refSignature;
			if (icompare(method, "PLAINTEXT") == 0)
			{
				refSignature = percentEncode(m_consumerSecret);
				refSignature += '&';
				refSignature += percentEncode(m_tokenSecret);
			}
			else if (icompare(method, "HMAC-SHA1") == 0)
			{
				CURI uriWithoutQuery(uri);
				uriWithoutQuery.setQuery("");
				uriWithoutQuery.setFragment("");
				refSignature = createSignature(request, uriWithoutQuery.toString(), params, nonce, timestamp);
			}
			else throw NotAuthenticatedException("Unsupported OAuth signature method", method);

			return refSignature == signature;			
		}
		else throw NotAuthenticatedException("No OAuth credentials found in Authorization header");
	}
	else throw NotAuthenticatedException("No Authorization header found");
}

void COAuth10Credentials::nonceAndTimestampForTesting(const std::string& nonce, const std::string& timestamp)
{
	m_nonce = nonce;
	m_timestamp = timestamp;
}

void COAuth10Credentials::signPlaintext(CHTTPRequest& request) const
{
	std::string signature(percentEncode(m_consumerSecret));
	signature += '&';
	signature += percentEncode(m_tokenSecret);

	std::string authorization(SCHEME);
	if (!m_realm.empty())
	{
		format(authorization, " realm=\"%s\",", m_realm);
	}
	format(authorization, " oauth_consumer_key=\"%s\"", percentEncode(m_consumerKey));
	format(authorization, ", oauth_signature=\"%s\"", percentEncode(signature));
	authorization += ", oauth_signature_method=\"PLAINTEXT\"";
	if (!m_token.empty())
	{
		format(authorization, ", oauth_token=\"%s\"", percentEncode(m_token));
	}
	if (!m_callback.empty())
	{
		format(authorization, ", oauth_callback=\"%s\"", percentEncode(m_callback));
	}
	authorization += ", oauth_version=\"1.0\"";

	request.set(CHTTPRequest::AUTHORIZATION, authorization);
}

void COAuth10Credentials::signHMACSHA1(CHTTPRequest& request, const std::string& uri, const CHTMLForm& params) const
{
	std::string nonce(m_nonce);
	if (nonce.empty())
	{
		nonce = createNonce();
	}
	std::string timestamp(m_timestamp);
	if (timestamp.empty())
	{
		timestamp = CNumberFormatter::format(CTimeStamp().epochTime());
	}
	std::string signature(createSignature(request, uri, params, nonce, timestamp));

	std::string authorization(SCHEME);
	if (!m_realm.empty())
	{
		format(authorization, " realm=\"%s\",", m_realm);
	}
	format(authorization, " oauth_consumer_key=\"%s\"", percentEncode(m_consumerKey));
	format(authorization, ", oauth_nonce=\"%s\"", percentEncode(nonce));
	format(authorization, ", oauth_signature=\"%s\"", percentEncode(signature));
	authorization += ", oauth_signature_method=\"HMAC-SHA1\"";
	format(authorization, ", oauth_timestamp=\"%s\"", timestamp);
	if (!m_token.empty())
	{
		format(authorization, ", oauth_token=\"%s\"", percentEncode(m_token));
	}
	if (!m_callback.empty())
	{
		format(authorization, ", oauth_callback=\"%s\"", percentEncode(m_callback));
	}
	authorization += ", oauth_version=\"1.0\"";

	request.set(CHTTPRequest::AUTHORIZATION, authorization);
}

std::string COAuth10Credentials::createNonce() const
{
	std::ostringstream base64Nonce;
	CBase64Encode base64Encoder(base64Nonce);
	CRandomInputStream randomStream;
	for (int i = 0; i < 32; i++)
	{
		base64Encoder.put(randomStream.get());
	}
	base64Encoder.close();
	std::string nonce = base64Nonce.str();
	return translate(nonce, "+/=", "");
}

std::string COAuth10Credentials::createSignature(const CHTTPRequest& request, const std::string& uri, const CHTMLForm& params, const std::string& nonce, const std::string& timestamp) const
{
	std::map<std::string, std::string> paramsMap;
	paramsMap["oauth_version"]          = "1.0";
	paramsMap["oauth_consumer_key"]     = percentEncode(m_consumerKey);
	paramsMap["oauth_nonce"]            = percentEncode(nonce);
	paramsMap["oauth_signature_method"] = "HMAC-SHA1";
	paramsMap["oauth_timestamp"]        = timestamp;
	if (!m_token.empty())
	{
		paramsMap["oauth_token"] = percentEncode(m_token);
	}
	if (!m_callback.empty())
	{
		paramsMap["oauth_callback"] = percentEncode(m_callback);
	}
	for (CHTMLForm::ConstIterator it = params.begin(); it != params.end(); ++it)
	{
		paramsMap[percentEncode(it->first)] = percentEncode(it->second);
	}

	std::string paramsString;
	for (std::map<std::string, std::string>::const_iterator it = paramsMap.begin(); it != paramsMap.end(); ++it)
	{
		if (it != paramsMap.begin()) paramsString += '&';
		paramsString += it->first;
		paramsString += "=";
		paramsString += it->second;
	}

	std::string signatureBase = request.getMethod();
	signatureBase += '&';
	signatureBase += percentEncode(uri);
	signatureBase += '&';
	signatureBase += percentEncode(paramsString);

	std::string signingKey;
	signingKey += percentEncode(m_consumerSecret);
	signingKey += '&';
	signingKey += percentEncode(m_tokenSecret);

	THMACEngine<CSHA1Engine> hmacEngine(signingKey);
	hmacEngine.update(signatureBase);
	CDigestEngine::Digest digest = hmacEngine.digest();
	std::ostringstream digestBase64;
	CBase64Encode base64Encoder(digestBase64);
	base64Encoder.write(reinterpret_cast<char*>(&digest[0]), digest.size());
	base64Encoder.close();
	return digestBase64.str();
}

std::string COAuth10Credentials::percentEncode(const std::string& str)
{
	std::string encoded;
	CURI::encode(str, "!?#/'\",;:$&()[]*+=@", encoded);
	return encoded;
}

NET_NAMESPACE_END
